package cn.xiaobai.tools;

import org.eclipse.sisu.space.asm.*;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * @Author yangdaji
 * @Date 2023/10/26
 * @PackageName cn.xiaobai.tools
 * @ClassName MainClassFinder
 */
public class MainClassFinder {

    private static final String DOT_CLASS = ".class";

    public static String findSingleMainClass(JarFile source, String annotationName) throws IOException {
        SingleMainClassCallback callback = new SingleMainClassCallback(annotationName);
        MainClassFinder.doWithMainClasses(source, callback);
        return callback.getMainClassName();
    }

    static <T> T doWithMainClasses(JarFile jarFile, MainClassCallback<T> callback)
            throws IOException {
        Enumeration<JarEntry> entries = jarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            String name = entry.getName();
            if (name.endsWith(DOT_CLASS)) {
                try (InputStream inputStream = new BufferedInputStream(jarFile.getInputStream(entry))) {
                    ClassReader classReader = new ClassReader(inputStream);
                    ClassDescriptor classDescriptor = new ClassDescriptor(ClassDescriptor.minMemoryCapacity);
                    classReader.accept(classDescriptor, ClassReader.SKIP_CODE);
                    if (classDescriptor.isMainMethodFound()) {
                        T result = callback.doWith(new MainClass(ClassConvertUtils.convertClassName(name), classDescriptor.getAnnotationNames()));
                        if (result != null) {
                            return result;
                        }
                    }
                }
            }
        }
        return null;
    }


    static final class ClassDescriptor extends ClassVisitor {


        private static final Type STRING_ARRAY_TYPE = Type.getType(String[].class);

        private static final Type MAIN_METHOD_TYPE = Type.getMethodType(Type.VOID_TYPE, STRING_ARRAY_TYPE);

        private static final String MAIN_METHOD_NAME = "main";
        private final Set<String> annotationNames = new LinkedHashSet<>();

        private boolean mainMethodFound;

        public static final int minMemoryCapacity = 256 * 1024;
        public static final int maxMemoryCapacity = 320 * 1024;

        ClassDescriptor(int capacity) {
            super(capacity == 0 ? minMemoryCapacity : capacity);
        }

        @Override
        public AnnotationVisitor visitAnnotation(String desc, boolean visible) {
            this.annotationNames.add(Type.getType(desc).getClassName());
            return null;
        }

        @Override
        public MethodVisitor visitMethod(int access, String name, String desc, String signature, String[] exceptions) {
            if (isAccess(access, Opcodes.ACC_PUBLIC, Opcodes.ACC_STATIC) && MAIN_METHOD_NAME.equals(name)
                    && MAIN_METHOD_TYPE.getDescriptor().equals(desc)) {
                this.mainMethodFound = true;
            }
            return null;
        }

        private boolean isAccess(int access, int... requiredOpsCodes) {
            for (int requiredOpsCode : requiredOpsCodes) {
                if ((access & requiredOpsCode) == 0) {
                    return false;
                }
            }
            return true;
        }

        public boolean isMainMethodFound() {
            return this.mainMethodFound;
        }

        Set<String> getAnnotationNames() {
            return this.annotationNames;
        }

    }

    interface MainClassCallback<T> {


        T doWith(MainClass mainClass);

    }

    static final class MainClass {

        private final String name;

        private final Set<String> annotationNames;

        MainClass(String name, Set<String> annotationNames) {
            this.name = name;
            this.annotationNames = Collections.unmodifiableSet(new HashSet<>(annotationNames));
        }

        String getName() {
            return this.name;
        }

        Set<String> getAnnotationNames() {
            return this.annotationNames;
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj) {
                return true;
            }
            if (obj == null) {
                return false;
            }
            if (getClass() != obj.getClass()) {
                return false;
            }
            MainClass other = (MainClass) obj;
            return this.name.equals(other.name);
        }

        @Override
        public int hashCode() {
            return this.name.hashCode();
        }

        @Override
        public String toString() {
            return this.name;
        }

    }

    private static final class SingleMainClassCallback implements MainClassCallback<Object> {

        private final Set<MainClass> mainClasses = new LinkedHashSet<>();

        private final String annotationName;

        private SingleMainClassCallback(String annotationName) {
            this.annotationName = annotationName;
        }

        @Override
        public Object doWith(MainClass mainClass) {
            this.mainClasses.add(mainClass);
            return null;
        }

        private String getMainClassName() {
            Set<MainClass> matchingMainClasses = new LinkedHashSet<>();
            if (this.annotationName != null) {
                for (MainClass mainClass : this.mainClasses) {
                    if (mainClass.getAnnotationNames().contains(this.annotationName)) {
                        matchingMainClasses.add(mainClass);
                    }
                }
            }
            if (matchingMainClasses.isEmpty()) {
                matchingMainClasses.addAll(this.mainClasses);
            }
            if (matchingMainClasses.size() > 1) {
                throw new IllegalStateException(
                        "Unable to find a single main class from the following candidates " + matchingMainClasses);
            }
            return (matchingMainClasses.isEmpty() ? null : matchingMainClasses.iterator().next().getName());
        }
    }
}
